# 集成分类器
class EnsembleNovelClassifier:
    def __init__(self):
        self.models = {
            'random_forest': RandomForestClassifier(
                n_estimators=200,
                max_depth=20,
                min_samples_split=5,
                min_samples_leaf=2,
                random_state=42
            ),
            'svm': SVC(
                kernel='linear',
                probability=True,
                random_state=42
            ),
            'logistic': LogisticRegression(
                C=1.0,
                max_iter=1000,
                random_state=42
            )
        }
        self.label_encoder = LabelEncoder()
        self.feature_engineer = FeatureEngineer()
        self.preprocessor = AdvancedTextPreprocessor()
        
    def prepare_data(self, texts, labels):
        """准备数据"""
        # 文本预处理
        processed_texts = [self.preprocessor.extract_features(text) for text in texts]
        
        # 特征工程
        features = self.feature_engineer.create_all_features(processed_texts)
        
        # 标签编码
        encoded_labels = self.label_encoder.fit_transform(labels)
        
        return features, encoded_labels
    
    def train(self, texts, labels, test_size=0.2):
        """训练模型"""
        # 准备数据
        X, y = self.prepare_data(texts, labels)
        
        # 划分训练测试集
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=test_size, random_state=42, stratify=y
        )
        
        # 训练各个模型
        trained_models = {}
        for name, model in self.models.items():
            print(f"训练 {name}...")
            model.fit(X_train, y_train)
            y_pred = model.predict(X_test)
            accuracy = accuracy_score(y_test, y_pred)
            print(f"{name} 准确率: {accuracy:.4f}")
            trained_models[name] = model
        
        self.models = trained_models
        return X_test, y_test
    
    def predict_ensemble(self, texts):
        """集成预测"""
        processed_texts = [self.preprocessor.extract_features(text) for text in texts]
        features = self.feature_engineer.create_all_features(processed_texts)
        
        # 获取各模型预测概率
        probabilities = []
        for model in self.models.values():
            proba = model.predict_proba(features)
            probabilities.append(proba)
        
        # 平均概率
        avg_proba = np.mean(probabilities, axis=0)
        predictions = np.argmax(avg_proba, axis=1)
        
        return self.label_encoder.inverse_transform(predictions), avg_proba
    
    def save_model(self, path):
        """保存模型"""
        model_data = {
            'models': self.models,
            'label_encoder': self.label_encoder,
            'feature_engineer': self.feature_engineer,
            'preprocessor': self.preprocessor
        }
        joblib.dump(model_data, path)
    
    def load_model(self, path):
        """加载模型"""
        model_data = joblib.load(path)
        self.models = model_data['models']
        self.label_encoder = model_data['label_encoder']
        self.feature_engineer = model_data['feature_engineer']
        self.preprocessor = model_data['preprocessor']